import mxnet as mx
from ..rpn.rpn import assign_anchor
import numpy as np
from mxnet.gluon.data import Dataset,DataLoader
class AnchorDataset(Dataset):
    def __init__(self, bbox_dataset, cfg=None, transform = None):
        super(AnchorDataset).__init__()
        self.cfg = cfg

        self.bbox_dataset = bbox_dataset

        self.transform = transform
    def __getitem__(self, idx):
        image, label = self.bbox_dataset[idx]
        label[:,4] += 1 # let 0 be background
        assert len(label) > 0
        assert np.all(label[:,4] <= self.cfg.dataset.NUM_CLASSES)
        if self.transform is not None:
            image,label = self.transform(image,label)
        h,w,nch = image.shape
        # assert w % 8 == 0, w % 8 == 0
        assert self.cfg.NETWORK_STRIDE[0] == self.cfg.NETWORK_STRIDE[1]
        feat_stride = self.cfg.NETWORK_STRIDE[0]
        rpn_batch = assign_anchor(feat_shape=[1, 1024,
                                          h // self.cfg.NETWORK_STRIDE[0],
                                          w // self.cfg.NETWORK_STRIDE[1]
                                          ],
                              gt_boxes = label[:,:5],
                              im_info = [[w,h,nch]],
                              cfg= self.cfg,
                              feat_stride=feat_stride,
                              scales = self.cfg.TRAIN.RPN_ANCHOR_SCALES,
                              ratios= self.cfg.TRAIN.RPN_ANCHOR_RATIOS
                              )
        img_float = image.astype(np.float32)
        mean = np.array([103.06,115.90, 123.15])[np.newaxis,np.newaxis].astype(np.float32)
        img_float = img_float - mean
        img_float = img_float[:,:,(2,1,0)]
        img_float = np.transpose(img_float,(2,0,1))
        
        im_info = np.array([img_float.shape[1],img_float.shape[2],1])
        gt_boxes = np.zeros(shape = (300,label.shape[1]))
        gt_boxes[:label.shape[0],:label.shape[1]] = label
        return image,img_float,im_info,rpn_batch["cls_label_onehot"][0],rpn_batch['cls_weight'][0],rpn_batch["bbox_target"][0].astype( np.float32),rpn_batch["bbox_weight"][0].astype( np.float32),gt_boxes.astype(np.float32)

    def __len__(self):
        return len(self.bbox_dataset)

